// Copyright [2022] <Tianyu Zhu, Shizhen Zhao, Zihang Ren>

#include <algorithm>
#include <cstdlib>  // Header file needed to use srand and rand
#include <iostream>

#include "pulse_plus.h"

PulsePlus::~PulsePlus() {
    if (cost_map_) {
        delete[] cost_map_;
        cost_map_ = nullptr;
    }
    if (min_cost_path_map_) {
        delete[] min_cost_path_map_;
        min_cost_path_map_ = nullptr;
    }
    if (delay_map_) {
        delete[] delay_map_;
        delay_map_ = nullptr;
    }
    if (visited_) {
        delete[] visited_;
        visited_ = nullptr;
    }
    if (info_) {
        delete[] info_;
        info_ = nullptr;
    }
}

void PulsePlus::SetupTopology(Graph* graph) {
    graph_ = graph;
    num_nodes_ = graph_->NodeSize();
    cost_map_ = new int[num_nodes_];
    min_cost_path_map_ = new Path[num_nodes_];
    delay_map_ = new int[num_nodes_];
    visited_ = new bool[num_nodes_];
    info_ = new PulseInfo[num_nodes_]();
    srand(0);
}

Path PulsePlus::FindPath(const Flow& flow) {
    // Initialization
    // do dijkstra from end to every other nodes,
    // from two perspectives--cost and delay
    // store the result of dijkstras in cost_map and delay_map
    flow_ = flow;
    clock_t start_time = clock();
    CostBasedDijkstra(flow_.to);
    DelayBasedDijkstra(flow_.to);
    clock_t end_time = clock();
    std::cout << "PulsePlus initialization time: "
              << double(end_time - start_time) / CLOCKS_PER_SEC * 1000
              << "(ms).\n";

    // Reset Pulse Algorithm
    int cnt = 0;
    for (int i = 0; i < num_nodes_; ++i) {
        visited_[i] = false;
        info_[i].Clear();
        min_cost_path_map_[i].clear();
    }
    NodeId start = flow_.from;
    visited_[start] = true;
    total_cost_ = kMaxValue;
    results_.clear();

    // Timer starts
    start_time = clock();
    // Start pulse algorithm
    Path path;
    path.path_info.push_back(start);
    path.cost = 0;
    path.delay = 0;
    Dfs(start, path);
    // Print time elapse
    end_time = clock();
    std::cout << "PulsePlus time: "
              << double(end_time - start_time) / CLOCKS_PER_SEC * 1000
              << "(ms).\n";
    return results_.back();
}

// dijkstra algorithm
void PulsePlus::CostBasedDijkstra(NodeId dst_node) {
    for (int i = 0; i < num_nodes_; ++i) {
        cost_map_[i] = kMaxValue;
    }
    cost_map_[dst_node] = 0;
    std::priority_queue<Node, std::vector<Node>> heap;
    std::vector<bool> visited(num_nodes_, false);
    heap.push(Node(dst_node, 0, -1, 0));
    while (!heap.empty()) {
        Node cur_node = heap.top();
        NodeId u = cur_node.id;
        heap.pop();
        if (visited[u]) continue;
        visited[u] = true;
        NodeId pre_id = cur_node.pre_id;
        int weight_u_to_dst = cur_node.distance;
        min_cost_path_map_[u].path_info.push_back(u);
        min_cost_path_map_[u].cost = weight_u_to_dst;
        if (pre_id != -1) {
            min_cost_path_map_[u].delay =
                min_cost_path_map_[pre_id].delay + cur_node.link_delay;
            for (NodeId node_id : min_cost_path_map_[pre_id].path_info) {
                min_cost_path_map_[u].path_info.push_back(node_id);
            }
        } else {
            min_cost_path_map_[u].delay = 0;
        }
        for (Link* link : graph_->GetIngressLinks(u)) {
            NodeId v = link->source_id;
            if (visited[v]) continue;
            int weight = link->cost;
            if (cost_map_[v] > weight_u_to_dst + weight) {
                cost_map_[v] = weight_u_to_dst + weight;
                heap.push(Node(v, cost_map_[v], u, link->delay));
            }
        }
    }
}

// dijkstra algorithm
void PulsePlus::DelayBasedDijkstra(NodeId dst_node) {
    for (int i = 0; i < num_nodes_; ++i) {
        delay_map_[i] = kMaxValue;
    }
    delay_map_[dst_node] = 0;
    std::priority_queue<Node, std::vector<Node>> heap;
    std::vector<bool> visited(num_nodes_, false);
    heap.push(Node(dst_node, 0, -1, 0));
    while (!heap.empty()) {
        NodeId u = heap.top().id;
        int weight_u_to_dst = heap.top().distance;
        heap.pop();
        if (visited[u]) continue;
        visited[u] = true;
        for (Link* link : graph_->GetIngressLinks(u)) {
            NodeId v = link->source_id;
            if (visited[v]) continue;
            int weight = link->delay;
            if (delay_map_[v] > weight_u_to_dst + weight) {
                delay_map_[v] = weight_u_to_dst + weight;
                // The last two fields are not used here.
                heap.push(Node(v, delay_map_[v], -1, 0));
            }
        }
    }
}

// prune accordding to bound and infeasibility
bool PulsePlus::check(NodeId u, int cost, int delay) {
    if (cost_map_[u] + cost >= total_cost_)
        return false;
    if (delay_map_[u] + delay > flow_.delay_ub)
        return false;
    return true;
}

// update the path node U records
// according to the rule mentioned in Chapter 4 of 2012.10 paper
bool PulsePlus::update(NodeId u, const Path& path, double p) {
    int delay = path.delay;
    auto it = info_[u].delay_to_cost_map.lower_bound(delay);
    if (it == info_[u].delay_to_cost_map.end()) {
        info_[u].delay_to_cost_map[delay] = path.cost;
        return true;
    }
    // The searched value is in the map
    if (it->first == delay) {
        if (it->second <= path.cost) {
            return false;
        } else {
            info_[u].delay_to_cost_map[delay] = path.cost;
            return true;
        }
    }
    if (it == info_[u].delay_to_cost_map.begin()) {
        info_[u].delay_to_cost_map[delay] = path.cost;
        return true;
    }
    if (it->second > path.cost) {
        info_[u].delay_to_cost_map[delay] = path.cost;
        return true;
    }
    int post_delay = it->first;
    --it;
    if (post_delay - it->first > flow_.delay_ub - flow_.delay_lb) {
        info_[u].delay_to_cost_map[delay] = path.cost;
        return true;
    }
    if (it->second > path.cost) {
        info_[u].delay_to_cost_map[delay] = path.cost;
        return true;
    }
    return false;
}

// dfs from start to end, use recursion
void PulsePlus::Dfs(NodeId u, Path& path) {
    if (u == flow_.to) {
        if (path.delay > flow_.delay_ub || path.delay < flow_.delay_lb) {
            return;
        }
        if (path.cost > total_cost_) {
            return;
        }
        // total_cost_ = cost * 0.8;
        total_cost_ = path.cost;
        results_.push_back(path);
        return;
    }
    // prune
    if (!check(u, path.cost, path.delay)) {
        return;
    }
    if (path.delay + min_cost_path_map_[u].delay <= flow_.delay_ub &&
        path.delay + min_cost_path_map_[u].delay >= flow_.delay_lb) {
        const Path& min_cost_path = min_cost_path_map_[u];
        bool loop = false;
        for (int i = 1; i < min_cost_path.path_info.size(); ++i) {
            if (visited_[min_cost_path.path_info[i]]) {
                // Find a loop
                loop = true;
                break;
            }
        }
        if (!loop) {
            total_cost_ = path.cost + min_cost_path.cost;
            results_.push_back(path);
            Path& last_path = results_.back();
            for (int i = 1; i < min_cost_path.path_info.size(); ++i) {
                last_path.path_info.push_back(min_cost_path.path_info[i]);
            }
            last_path.cost += min_cost_path.cost;
            last_path.delay += min_cost_path.delay;
            return;
        }
    }
    double p = static_cast<double>(rand() % 100) / 100;
    // Use dominance check to prune branches only if there
    // is no delay lower bound.
    // if (!update(u, path, p)) {
    //     return;
    // }
    for (Link* link : graph_->GetEgressLinks(u)) {
        NodeId ter = link->ter_id;
        if (visited_[ter]) {
            continue;
        }
        visited_[ter] = true;
        path.cost += link->cost;
        path.delay += link->delay;
        path.path_info.push_back(ter);
        Dfs(ter, path);
        path.cost -= link->cost;
        path.delay -= link->delay;
        path.path_info.pop_back();
        visited_[ter] = false;
    }
}

void PulsePlus::print() const {
    std::vector<Path> buf;
    for (int i = 0; i < results_.size(); ++i) {
        if (results_[i].delay >= flow_.delay_lb &&
            results_[i].delay <= flow_.delay_ub) {
            buf.push_back(results_[i]);
        }
    }
    results_.back().Print();
}
